
#include "Common.hpp"
using namespace std;

namespace {

//-----------------------------------------------------------------------------
const char* nVertexProgramSrc =
"!!NVvp5.0"
NL
NL "main:"
NL "    MOV.F32 result.position, vertex.attrib[0];"
NL "    RET;"
NL "END";

const char* nFragmentProgramSrc =
"!!NVfp5.0"
NL
NL "main:"
NL "    TEX result.color.rgb, fragment.position, texture[0], RECT;"
NL "    MOV.F32 result.color.a, 1.0;"
NL "    RET;"
NL "END";

const char* nRaytraceModuleSrc =
NL ".version 2.3"
NL ".target sm_20"
NL ".address_size 32"

NL ".global .surfref _fb;"
NL ".const .v4 .f32 _eye = { 0.0, 0.0, 5.0, 1.0 };"
NL ".const .v4 .f32 _view[4] = {"
NL "    { 1.0, 0.0, 0.0, 0.0 },"
NL "    { 0.0, 1.0, 0.0, 0.0 },"
NL "    { 0.0, 0.0, 1.0, 0.0 },"
NL "    { 0.0, 0.0, 0.0, 1.0 }"
NL "};"

#define OBJECT_COUNT "3"
#define OBJECT_SIZE "32"
#define OBJECT_TYPE_SPHERE "1.0"
#define OBJECT_TYPE_PLANE "2.0"

NL ".const .v4 .f32 _object["OBJECT_COUNT"][2] = {"
NL "    { { 0.7, 0.7, 0.0, "OBJECT_TYPE_SPHERE" }, { 1.0, -0.8, -3.0, 2.5 } },"
NL "    { { 0.7, 0.0, 0.1, "OBJECT_TYPE_SPHERE" }, { -5.5, -0.5, -7.0, 2.0 } },"
NL "    { { 0.4, 0.3, 0.3, "OBJECT_TYPE_PLANE" }, { 0.0, 1.0, 0.0, 4.4 } }"
NL "};"

#define LIGHT_COUNT "3"
#define LIGHT_SIZE "32"

NL ".const .v4 .f32 _light["LIGHT_COUNT"][2] = {"
NL "    { { 0.0, 5.0, -5.0, 1.0 }, { 0.6, 0.6, 0.6, 1.0 } },"
NL "    { { 2.0, 5.0, -1.0, 1.0 }, { 0.7, 0.7, 0.9, 1.0 } },"
NL "    { { 0.0, 15.0, 0.0, 1.0 }, { 0.8, 0.8, 0.8, 1.0 } }"
NL "};"

NL ".func (.reg .f32 dist) sphere_isect(.reg .v4 .f32 ro, .reg .v4 .f32 rd, .reg .v4 .f32 sphere) {"
NL "    .reg .v4 .f32 l;"
NL "    .reg .f32 s, l2, m2, q;"
NL "    .reg .pred p;"
        ptx_op_v3_v3("sub.f32", "l", "sphere", "ro")
        ptx_dot_v3(".f32", "s", "l", "rd")
        ptx_dot_v3(".f32", "l2", "l", "l")
NL "    mul.f32 sphere.w, sphere.w, sphere.w;"
NL "    neg.f32 q, s;"
NL "    fma.rn.f32 m2, q, s, l2;"
NL "    sub.f32 q, sphere.w, m2;"
NL "    sqrt.approx.f32 q, q;"
NL "    setp.gt.f32 p, l2, sphere.w;"
NL "@p  sub.f32 dist, s, q;"
NL "@!p add.f32 dist, s, q;"
NL "    ret;"
NL "}"

NL ".func (.reg .v4 .f32 normal) sphere_normal(.reg .v4 .f32 sphere, .reg .v4 .f32 isect_pt) {"
        ptx_op_v3_v3("sub.f32", "normal", "isect_pt", "sphere")
NL "    .reg .f32 f;"
NL "    rcp.approx.f32 f, sphere.w;"
        ptx_op_v3_s("mul.f32", "normal", "normal", "f")
NL "    ret;"
NL "}"

NL ".func (.reg .f32 dist) plane_isect(.reg .v4 .f32 ro, .reg .v4 .f32 rd, .reg .v4 .f32 plane) {"
NL "    .reg .f32 dp;"
        ptx_dot_v3(".f32", "dp", "rd", "plane")
        ptx_dot_v3(".f32", "dist", "ro", "plane")
NL "    add.f32 dist, dist, plane.w;"
NL "    div.approx.f32 dist, dist, dp;"
NL "    neg.f32 dist, dist;"
NL "    ret;"
NL "}"

NL ".func (.param .b8 res[8]) object_isect(.reg .v4 .f32 ro, .reg .v4 .f32 rd) {"
NL "    .reg .f32 dist, min_dist;"
NL "    .reg .s32 i, min_obj;"
NL "    .reg .pred p;"
NL "    .reg .v4 .f32 obj_param0, obj_param1;"
NL "    .reg .u32 ptr;"

NL "    mov.u32 ptr, _object;"
NL "    mov.f32 min_dist, 10000.0;"
NL "    mov.s32 min_obj, -1;"
NL "    mov.s32 i, 0;"

NL "object_loop:" // loop over each object in the scene
NL "    ld.const.v4.f32 obj_param0, [ptr];"
NL "    ld.const.v4.f32 obj_param1, [ptr + 16];"

NL "    setp.eq.f32 p, obj_param0.w, "OBJECT_TYPE_SPHERE";"
NL "    @p call (dist), sphere_isect, (ro, rd, obj_param1);"

NL "    setp.eq.f32 p, obj_param0.w, "OBJECT_TYPE_PLANE";"
NL "    @p call (dist), plane_isect, (ro, rd, obj_param1);"

NL "    setp.gt.f32 p, dist, 0.0;"
NL "    setp.lt.and.f32 p, dist, min_dist, p;"
NL "    @p mov.s32 min_obj, i;"
NL "    @p mov.f32 min_dist, dist;"

NL "    add.u32 ptr, ptr, "OBJECT_SIZE";"
NL "    add.s32 i, i, 1;"
NL "    setp.lt.s32 p, i, "OBJECT_COUNT";"
NL "    @p bra.uni object_loop;" // object loop end

NL "    st.param.s32 [res], min_obj;"
NL "    st.param.f32 [res + 4], min_dist;"
NL "    ret;"
NL "}"

NL ".func (.reg .v4 .f32 color) shade(.reg .v4 .f32 ro, .reg .v4 .f32 rd, .reg .v4 .f32 c, .reg .v4 .f32 p, .reg .v4 .f32 n) {"
NL "    .reg .s32 i;"
NL "    .reg .u32 address;"
NL "    .reg .pred q;"
NL "    .reg .v4 .f32 l_pos, l_col, l_vec, l_vec_n;"

NL "    mov.s32 i, 0;"
NL "    mov.u32 address, _light;"
NL "    mov.v4.f32 color, { 0.0, 0.0, 0.0, 1.0 };"

NL "light_loop:"
NL "    ld.const.v4.f32 l_pos, [address];"
NL "    ld.const.v4.f32 l_col, [address + 16];"

        ptx_op_v3_v3("sub.f32", "l_vec", "l_pos", "p")
        ptx_dot_v3(".f32", "l_vec.w", "l_vec", "l_vec")
NL      "rsqrt.approx.f32 l_vec.w, l_vec.w;"
        ptx_op_v3_s("mul.f32", "l_vec_n", "l_vec", "l_vec.w")

        ptx_op_v3_s_v3("fma.rn.f32", "l_pos", "l_vec_n", "0.001", "p")

NL "    .reg .f32 min_dist;"
NL "    .reg .s32 min_obj;"
NL "    .param .b8 isect_info[8];"
NL "    call.uni (isect_info), object_isect, (l_pos, l_vec_n);"
NL "    ld.param.s32 min_obj, [isect_info];"
NL "    ld.param.f32 min_dist, [isect_info + 4];"

NL "    .reg .f32 n_dot_l;"
        ptx_dot_sat_v3(".f32", "n_dot_l", "l_vec_n", "n")

NL "    rcp.approx.f32 l_vec.w, l_vec.w;"
NL "    setp.ge.f32 q, l_vec.w, min_dist;"
NL "    @q mov.f32 n_dot_l, 0.0;"

        ptx_op_v3_v3("mul.f32", "c", "c", "l_col")
        ptx_op_v3_s_v3("fma.rn.f32", "color", "c", "n_dot_l", "color")

NL "    add.u32 address, address, "LIGHT_SIZE";"
NL "    add.s32 i, i, 1;"
NL "    setp.lt.s32 q, i, "LIGHT_COUNT";"
NL "    @q bra.uni light_loop;"
NL "    ret;"
NL "}"

NL ".func (.reg .v4 .f32 color) trace(.reg .v4 .f32 ro, .reg .v4 .f32 rd) {"
NL "    .reg .pred p;"

NL "    .reg .f32 min_dist;"
NL "    .reg .s32 min_obj;"
NL "    .param .b8 isect_info[8];"
NL "    call.uni (isect_info), object_isect, (ro, rd);"
NL "    ld.param.s32 min_obj, [isect_info];"
NL "    ld.param.f32 min_dist, [isect_info + 4];"

NL "    setp.lt.s32 p, min_obj, 0;"
NL "    @p mov.v4.f32 color, { 0.0, 0.0, 0.0, 1.0 };"
NL "    @p ret;"

NL "    .reg .u32 ptr;"
NL "    .reg .v4 .f32 obj_param0, obj_param1;"
NL "    mov.u32 ptr, _object;"
NL "    mad.lo.u32 ptr, min_obj, "OBJECT_SIZE", ptr;"
NL "    ld.const.v4.f32 obj_param0, [ptr];"
NL "    ld.const.v4.f32 obj_param1, [ptr + 16];"

NL "    .reg .v4 .f32 isect_p, isect_n;"
        ptx_op_v3_s_v3("fma.rn.f32", "isect_p", "rd", "min_dist", "ro")
NL "    setp.eq.f32 p, obj_param0.w, "OBJECT_TYPE_SPHERE";"
NL "    @p call (isect_n), sphere_normal, (obj_param1, isect_p);"
NL "    setp.eq.f32 p, obj_param0.w, "OBJECT_TYPE_PLANE";"
NL "    @p mov.v4.f32 isect_n, obj_param1;"

NL "    call.uni (color), shade, (ro, rd, obj_param0, isect_p, isect_n);"
NL "    ret;"
NL "}"

NL ".entry raytrace() {"
NL "    .reg .u32 u<3>;"
        // compute x coordinate: x = %ctaid.x * %ntid.x + %tid.x;"
NL "    .reg .u32 x;"
NL "    mov.u32 u0, %ctaid.x;"
NL "    mov.u32 u1, %ntid.x;"
NL "    mov.u32 u2, %tid.x;"
NL "    mad.lo.u32 x, u0, u1, u2;"
        // compute y coordinate: y = %ctaid.y * %ntid.y + %tid.y;"
NL "    .reg .u32 y;"
NL "    mov.u32 u0, %ctaid.y;"
NL "    mov.u32 u1, %ntid.y;"
NL "    mov.u32 u2, %tid.y;"
NL "    mad.lo.u32 y, u0, u1, u2;"
        // compute ray origin"
NL "    .reg .v4 .f32 ro, view<4>;"
NL "    .reg .f32 f0;"
NL "    cvt.rn.f32.u32 f0, x;"
NL "    fma.rn.f32 ro.x, f0, 1.0 / 1024.0, -0.5;"
NL "    mul.f32 ro.x, ro.x, 8.0;"
NL "    cvt.rn.f32.u32 f0, y;"
NL "    fma.rn.f32 ro.y, f0, 1.0 / 1024.0, -0.5;"
NL "    mul.f32 ro.y, ro.y, 8.0;"
NL "    mov.f32 ro.z, 0.0;"
NL "    mov.f32 ro.w, 1.0;"
NL "    ld.const.v4.f32 view0, [_view];"
NL "    ld.const.v4.f32 view1, [_view + 16];"
NL "    ld.const.v4.f32 view2, [_view + 32];"
NL "    ld.const.v4.f32 view3, [_view + 48];"
        ptx_mul_m4_v4(".f32", "ro", "view", "ro")
        // compute ray direction: rd = normalize(ro - eye);"
NL "    .reg .v4 .f32 rd, eye;"
NL "    ld.const.v4.f32 eye, [_eye];"
        ptx_mul_m4_v4(".f32", "eye", "view", "eye")
        ptx_op_v4_v4("sub.f32", "rd", "ro", "eye")
        ptx_dot_v3(".f32", "f0", "rd", "rd")
NL "    rsqrt.approx.f32 f0, f0;"
        ptx_op_v3_s("mul.f32", "rd", "rd", "f0")
        // trace
NL "    .reg .v4 .f32 color;"
NL "    call.uni (color), trace, (ro, rd);"
NL "    mul.lo.u32 x, x, 16;"
NL "    sust.b.2d.v4.b32.trap [_fb, { x, y }], color;"
NL "    ret;"
NL "}";
//-----------------------------------------------------------------------------
Window* nWindow;
ArcBall nArcBall;

SDL_GLContext nGLContext;
GLuint nProgram[2];
GLuint nVBO;
GLuint64 nVBOAddress;
GLuint nGLTexture;

CUcontext nCUContext;
CUmodule nCUModule;
CUfunction nCUFunction;
CUgraphicsResource nCUTexture;
CUsurfref nCUSurface;
CUdeviceptr nViewPtr;

glm::vec3 nCameraPos;
float nCameraRotX, nCameraRotY;

Uint32 nLastTime;

void UpdateCamera(float dt);
glm::vec3 ComputeDirection(float rotX, float rotY);
//-----------------------------------------------------------------------------
void Setup() {
    // window
    nWindow = new Window("Raytracing", 1024, 1024);

    // OpenGL context
    nGLContext = SDL_GL_CreateContext(nWindow->Get());
    if (!nGLContext) {
        throw runtime_error("[Setup] SDL_GL_CreateContext failed.");
    }

    SDL_GL_SetSwapInterval(0);

    if (glewInit() != GLEW_OK) {
        throw runtime_error("[Setup] glewInit failed.");
    }
    if (!GLEW_NV_gpu_program4) {
        throw runtime_error("[Setup] GL_NV_gpu_program4 extension is not supported.");
    }

    // programs
    glGenProgramsARB(2, nProgram);

    glBindProgramARB(GL_VERTEX_PROGRAM_ARB, nProgram[0]);
    LoadGLProgram(GL_VERTEX_PROGRAM_ARB, nVertexProgramSrc);

    glBindProgramARB(GL_FRAGMENT_PROGRAM_ARB, nProgram[1]);
    LoadGLProgram(GL_FRAGMENT_PROGRAM_ARB, nFragmentProgramSrc);

    // vbo
    const float vertices[] = { -1, -1, -1, 1, 1, -1, 1, 1 };
    glGenBuffers(1, &nVBO);
    glBindBuffer(GL_ARRAY_BUFFER, nVBO);
    glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);
    glGetBufferParameterui64vNV(GL_ARRAY_BUFFER, GL_BUFFER_GPU_ADDRESS_NV, &nVBOAddress);
    glMakeBufferResidentNV(GL_ARRAY_BUFFER, GL_READ_ONLY);

    glGenTextures(1, &nGLTexture);
    glBindTexture(GL_TEXTURE_RECTANGLE, nGLTexture);
    glTexParameteri(GL_TEXTURE_RECTANGLE, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_RECTANGLE, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexImage2D(GL_TEXTURE_RECTANGLE, 0, GL_RGBA32F, nWindow->Width(), nWindow->Height(), 0, GL_RGBA, GL_FLOAT, NULL);

    // GL state
    glEnable(GL_VERTEX_PROGRAM_ARB);
    glEnable(GL_FRAGMENT_PROGRAM_ARB);

    // CUDA
    if (cuInit(0) != CUDA_SUCCESS) {
        throw runtime_error("[Setup] cuInit failed.");
    }

    CUdevice device;
    cuDeviceGet(&device, 0);
    if (cuGLCtxCreate(&nCUContext, 0, device) != CUDA_SUCCESS) {
        throw runtime_error("[Setup] cuCtxCreate failed.");
    }

    cout << nRaytraceModuleSrc << endl;

    nCUModule = LoadPTXModule(nRaytraceModuleSrc);

    if (cuModuleGetFunction(&nCUFunction, nCUModule, "raytrace") != CUDA_SUCCESS) {
        throw runtime_error("[Setup] cuModuleGetFunction failed.");
    }

    if (cuGraphicsGLRegisterImage(&nCUTexture, nGLTexture, GL_TEXTURE_RECTANGLE,
        CU_GRAPHICS_REGISTER_FLAGS_WRITE_DISCARD | CU_GRAPHICS_REGISTER_FLAGS_SURFACE_LDST) != CUDA_SUCCESS) {
        throw runtime_error("[Setup] cuGraphicsGLRegisterBuffer failed.");
    }

    if (cuModuleGetSurfRef(&nCUSurface, nCUModule, "_fb") != CUDA_SUCCESS) {
        throw runtime_error("[Setup] cuModuleGetSurfRef failed.");
    }

    // init GUI
    TwInit(TW_OPENGL, NULL);
    TwWindowSize(nWindow->Width(), nWindow->Height());

    TwBar* bar = TwNewBar("Usage");
    TwDefine("Usage color='0 0 0' alpha=128 position='10 10' size='240 50'");
    TwDefine("Usage fontresizable=false resizable=false movable=false");
    TwAddButton(bar, "usage0", NULL, NULL, "label='- mouse with LB down to rotate'");

    nArcBall.SetWindowSize(nWindow->Width(), nWindow->Height());

    nCameraPos = glm::vec3(0.0f, 1.0f, 7.0f);
    nCameraRotX = nCameraRotY = 0;

    nLastTime = SDL_GetTicks();

    if (cuModuleGetGlobal(&nViewPtr, NULL, nCUModule, "_view") != CUDA_SUCCESS) {
        throw runtime_error("[Setup] cuModuleGetGlobal failed.");
    }
}
//-----------------------------------------------------------------------------
bool ProcessEvents() {
    SDL_Event evt;
    while (SDL_PollEvent(&evt)) {
        if (evt.type == SDL_QUIT) {
            return false;
        }

        if (ProcessCommonEvent(evt) || ProcessGUIEvent(evt)) {
            continue;
        }

        switch (evt.type) {
        case SDL_MOUSEBUTTONDOWN:
            if (evt.button.button == SDL_BUTTON_LEFT) {
                nArcBall.Begin(evt.button.x, evt.button.y);
                return true;
            }
            break;

        case SDL_MOUSEBUTTONUP:
            if (evt.button.button == SDL_BUTTON_LEFT) {
                nArcBall.End();
                return true;
            }
            break;

        case SDL_MOUSEMOTION:
            nArcBall.Drag(evt.motion.x, evt.motion.y);
            break;
        }
    }

    return true;
}
//-----------------------------------------------------------------------------
void Update() {
    const Uint32 now = SDL_GetTicks();
    const float dt = (now - nLastTime) * 0.001f;
    nLastTime = now;
    const float currentTime = now * 0.001f;

    UpdateCamera(dt);

    glm::mat4 view = glm::lookAt(nCameraPos, nCameraPos + ComputeDirection(nCameraRotX, nCameraRotY),
        glm::vec3(0.0f, 1.0f, 0.0f));
    view = glm::inverse(view);
    view = glm::transpose(view);

    cuMemcpyHtoD(nViewPtr, &view[0][0], 64);

    // launch CUDA kernel
    cuGraphicsMapResources(1, &nCUTexture, NULL);

    CUarray arr;
    cuGraphicsSubResourceGetMappedArray(&arr, nCUTexture, 0, 0);
    cuSurfRefSetArray(nCUSurface, arr, 0);
    cuLaunchKernel(nCUFunction, 64, 64, 1, 16, 16, 1, 0, NULL, NULL, NULL);

    cuGraphicsUnmapResources(1, &nCUTexture, NULL);

    // draw
    glBindProgramARB(GL_VERTEX_PROGRAM_ARB, nProgram[0]);
    glBindProgramARB(GL_FRAGMENT_PROGRAM_ARB, nProgram[1]);

    glBindMultiTextureEXT(GL_TEXTURE0, GL_TEXTURE_RECTANGLE, nGLTexture);

    glEnableClientState(GL_VERTEX_ATTRIB_ARRAY_UNIFIED_NV);
    glEnableVertexAttribArray(0);
    glVertexAttribFormatNV(0, 2, GL_FLOAT, GL_FALSE, 2 * sizeof(float));
    glBufferAddressRangeNV(GL_VERTEX_ATTRIB_ARRAY_ADDRESS_NV, 0, nVBOAddress, 4 * 2 * sizeof(float));

    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);

    // draw GUI
    glDisableClientState(GL_VERTEX_ATTRIB_ARRAY_UNIFIED_NV);
    glDisableClientState(GL_ELEMENT_ARRAY_UNIFIED_NV);
    for (GLuint i = 0; i < 16; ++i) {
        glDisableVertexAttribArray(i);
    }
    TwDraw();

    nWindow->Swap();
}
//-----------------------------------------------------------------------------
void Run() {
    Setup();

    while (ProcessEvents()) {
        Update();
    }
}
//-----------------------------------------------------------------------------
void Cleanup() {
    if (nCUTexture) {
        cuGraphicsUnregisterResource(nCUTexture);
    }
    if (nCUModule) {
        cuModuleUnload(nCUModule);
    }
    if (nCUContext) {
        cuCtxDestroy(nCUContext);
    }

    glDeleteBuffers(1, &nVBO);
    glDeleteProgramsARB(2, nProgram);
    glDeleteTextures(1, &nGLTexture);

    TwTerminate();

    SDL_GL_DeleteContext(nGLContext);
    delete nWindow;
}
//-----------------------------------------------------------------------------
glm::vec3 ComputeDirection(float rotX, float rotY) {
    const glm::vec3 rotDir = glm::mat3(glm::rotate(glm::degrees(rotY), 0.0f, 1.0f, 0.0f)) * glm::vec3(0.0f, 0.0f, -1.0f);
    const glm::vec3 rotXAxis = -glm::cross(glm::vec3(0.0f, 1.0f, 0.0f), rotDir);
    return glm::mat3(glm::rotate(glm::degrees(rotX), rotXAxis)) * rotDir;
}
//-----------------------------------------------------------------------------
void UpdateCamera(float dt) {
    const float moveScale = 10 * dt;
    const float rotScale = 2 * dt;
    const glm::vec3 forward = ComputeDirection(nCameraRotX, nCameraRotY);
    const glm::vec3 right = -glm::cross(glm::vec3(0.0f, 1.0f, 0.0f), ComputeDirection(nCameraRotX, nCameraRotY));

    const Uint8* keyState = SDL_GetKeyboardState(NULL);

    if (keyState[SDL_SCANCODE_W]) {
        nCameraPos += forward * moveScale;
    }
    if (keyState[SDL_SCANCODE_S]) {
        nCameraPos -= forward * moveScale;
    }
    if (keyState[SDL_SCANCODE_D]) {
        nCameraPos += right * moveScale;
    }
    if (keyState[SDL_SCANCODE_A]) {
        nCameraPos -= right * moveScale;
    }

    int rx, ry;
    Uint8 buttonState = SDL_GetRelativeMouseState(&rx, &ry);
    if (buttonState & SDL_BUTTON_LMASK) {
        nCameraRotX -= ry * 0.005f;
        nCameraRotY -= rx * 0.005f;
    }
}
//-----------------------------------------------------------------------------

} // namespace

//-----------------------------------------------------------------------------
int main(int argc, char** argv) {
    if (SDL_Init(SDL_INIT_TIMER | SDL_INIT_VIDEO) < 0) {
        cout << "[main] SDL_Init failed." << endl;
        return 0;
    }

    try {
        Run();
    }
    catch (const exception& ex) {
        cout << ex.what() << endl;
    }
    catch (...) {
        cout << "Undefined exception." << endl;
    }

    Cleanup();
    SDL_Quit();

    return 0;
}
//-----------------------------------------------------------------------------
